import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import fetch_openml
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE, Isomap, LocallyLinearEmbedding, SpectralEmbedding
from sklearn.preprocessing import StandardScaler
import time
import os
import shutil
from warnings import simplefilter

# 忽略特定警告
simplefilter("ignore", category=RuntimeWarning)

# 设置缓存路径（可选）
# cache_path = os.path.join(os.path.expanduser('~'), 'sklearn_datasets')
# os.makedirs(cache_path, exist_ok=True)

print("正在加载MNIST数据集...")

# 方法1：尝试清除缓存（如果遇到问题）
# 注意：这会删除所有缓存数据集，谨慎使用
# from sklearn.datasets import clear_data_home
# clear_data_home()

# 加载MNIST数据集
try:
    # 使用更稳定的加载方式
    mnist = fetch_openml('mnist_784', version=1, as_frame=False, parser='auto')  # , data_home=cache_path)
    X = mnist.data
    y = mnist.target.astype(int)
    print("数据集加载成功！")
except Exception as e:
    print(f"加载失败: {e}")
    # 尝试替代加载方法
    from tensorflow.keras.datasets import mnist as keras_mnist

    (X_train, y_train), (X_test, y_test) = keras_mnist.load_data()
    X = np.vstack([X_train.reshape(-1, 784), X_test.reshape(-1, 784)])
    y = np.hstack([y_train, y_test])
    print("使用Keras备用数据集加载成功")

# 使用前5000个样本加速计算
sample_size = 5000
X = X[:sample_size]
y = y[:sample_size]

# 数据标准化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 可视化原始图像
print("\n显示原始图像...")
plt.figure(figsize=(12, 5))
for i in range(15):  # 显示更多样本
    plt.subplot(3, 5, i + 1)
    plt.imshow(X[i].reshape(28, 28), cmap='gray')
    plt.title(f"Label: {y[i]}")
    plt.axis('off')
plt.suptitle('MNIST原始图像示例 (15个样本)', fontsize=16)
plt.tight_layout()
plt.savefig('mnist_samples.png', dpi=150)
plt.show()

# PCA降维与可视化
print("\n进行PCA降维...")
start_time = time.time()
pca = PCA(n_components=2)
X_pca = pca.fit_transform(X_scaled)
pca_time = time.time() - start_time
print(f"PCA完成，耗时: {pca_time:.2f}秒")

plt.figure(figsize=(10, 8))
plt.scatter(X_pca[:, 0], X_pca[:, 1], c=y, cmap='tab10', alpha=0.7, s=15)
plt.colorbar(label='数字类别')
plt.title(f'PCA降维可视化 (解释方差: {pca.explained_variance_ratio_.sum():.2%})', fontsize=16)
plt.xlabel('主成分1')
plt.ylabel('主成分2')
plt.grid(alpha=0.2)
plt.savefig('pca_visualization.png', dpi=150)
plt.show()

# 定义降维方法列表
methods = [
    ("Isomap", Isomap(n_components=2, n_neighbors=10)),
    ("LLE", LocallyLinearEmbedding(n_components=2, n_neighbors=10, method='standard')),
    ("LE", SpectralEmbedding(n_components=2, n_neighbors=10, affinity='nearest_neighbors')),
    ("T-SNE", TSNE(n_components=2, perplexity=30, n_iter=1000, random_state=42))
]

# 执行各种降维算法并可视化
results = []
for name, model in methods:
    print(f"\n正在进行{name}降维...")
    start_time = time.time()

    try:
        X_trans = model.fit_transform(X_scaled)
        elapsed = time.time() - start_time
        results.append((name, X_trans, elapsed))

        plt.figure(figsize=(10, 8))
        plt.scatter(X_trans[:, 0], X_trans[:, 1], c=y, cmap='tab10', alpha=0.7, s=15)
        plt.colorbar(label='数字类别')
        plt.title(f'{name}降维可视化 (时间: {elapsed:.1f}秒)', fontsize=16)
        plt.xlabel('维度1')
        plt.ylabel('维度2')
        plt.grid(alpha=0.2)
        plt.savefig(f'{name.lower()}_visualization.png', dpi=150)
        plt.show()

        print(f"{name}完成，耗时: {elapsed:.2f}秒")

    except Exception as e:
        print(f"错误: {name}执行失败 - {str(e)}")

# 综合比较所有方法
plt.figure(figsize=(15, 12))
for i, (name, X_trans, elapsed) in enumerate(results):
    plt.subplot(2, 2, i + 1)
    plt.scatter(X_trans[:, 0], X_trans[:, 1], c=y, cmap='tab10', alpha=0.6, s=10)
    plt.title(f'{name} (时间: {elapsed:.1f}秒)', fontsize=14)
    plt.xticks([])
    plt.yticks([])

plt.suptitle('MNIST降维方法比较 (5000个样本)', fontsize=18)
plt.tight_layout()
plt.savefig('dimension_reduction_comparison.png', dpi=150)
plt.show()

# 打印解释方差比（仅PCA）
print("\nPCA解释方差比:", pca.explained_variance_ratio_)